#!/usr/bin/env python3

import argparse
import os
import sys
from time import sleep
import grpc
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))


def writeIpv4Rules(switch, dst_ip_addr, dst_eth_addr, port):

    table_entry = p4runtime_lib.helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": (dst_ip_addr, 32)},
        action_name="MyIngress.ipv4_forward",
        action_params={"dstAddr": dst_eth_addr, "port": port})
    switch.WriteTableEntry(table_entry)
    print(f"Installed IPv4 rule on {switch.name}: {dst_ip_addr} -> port {port}")


def writeArpRules(switch, dst_ip_addr, dst_eth_addr, port):

    table_entry = p4runtime_lib.helper.buildTableEntry(
        table_name="MyIngress.arp_table",
        match_fields={"hdr.arp.tpa": dst_ip_addr},
        action_name="MyIngress.arp_reply",
        action_params={"dstAddr": dst_eth_addr, "port": port})
    switch.WriteTableEntry(table_entry)
    print(f"Installed ARP rule on {switch.name}: {dst_ip_addr} -> port {port}")


def writeIpv6Rules(switch, dst_ip_addr, prefix_length, dst_eth_addr, port):

    table_entry = p4runtime_lib.helper.buildTableEntry(
        table_name="MyIngress.ipv6_lpm",
        match_fields={"hdr.ipv6.dstAddr": (dst_ip_addr, prefix_length)},
        action_name="MyIngress.ipv6_forward",
        action_params={"dstAddr": dst_eth_addr, "port": port})
    switch.WriteTableEntry(table_entry)
    print(f"Installed IPv6 rule on {switch.name}: {dst_ip_addr}/{prefix_length} -> port {port}")



def writeIcmpEncapsulationRule(switch, is_ipv4_encap, src_ipv4, dst_ipv4,
                               outer_dst_mac, outer_src_ip, outer_dst_ip, port):

    action_name = "MyIngress.encapsulate_icmp_ipv4" if is_ipv4_encap else "MyIngress.encapsulate_icmp_ipv6"
    match_fields = {
        "hdr.icmp.isValid": 1,
        "hdr.ipv4.srcAddr": (src_ipv4, 32),
        "hdr.ipv4.dstAddr": (dst_ipv4, 32)
    }
    action_params = {
        "dstAddr": outer_dst_mac,
        "port": port,
    }

    if is_ipv4_encap:
        action_params["outer_src_ip"] = outer_src_ip
        action_params["outer_dst_ip"] = outer_dst_ip
    else:
        action_params["outer_src_ipv6"] = outer_src_ip
        action_params["outer_dst_ipv6"] = outer_dst_ip

    table_entry = p4runtime_lib.helper.buildTableEntry(
        table_name="MyIngress.icmp_encapsulation",
        match_fields=match_fields,
        action_name=action_name,
        action_params=action_params)
    switch.WriteTableEntry(table_entry)
    encap_type = "IPv4" if is_ipv4_encap else "IPv6"
    print(f"Installed ICMP {encap_type} encapsulation rule on {switch.name}: "
          f"{src_ipv4} -> {dst_ipv4} -> port {port}")


def main(p4info_file_path, bmv2_file_path):
    p4info = p4runtime_lib.helper.read_p4info_fromfile(p4info_file_path)
    try:

        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')
        s4 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s4', address='127.0.0.1:50054', device_id=3,
            proto_dump_file='logs/s4-p4runtime-requests.txt')

        for sw in [s1, s2, s3, s4]:
            sw.MasterArbitrationUpdate()
            sw.SetForwardingPipelineConfig(
                p4info=None,
                bmv2_json_file_path=bmv2_file_path)
            print(f"Installed P4 program on {sw.name}")

        writeIpv4Rules(s1, "10.0.2.2", "08:00:00:00:02:22", 2)
        writeIpv4Rules(s2, "10.0.2.2", "08:00:00:00:02:22", 2)
        writeIpv4Rules(s4, "10.0.2.2", "08:00:00:00:02:22", 2)
        writeIpv4Rules(s4, "10.0.1.1", "08:00:00:00:01:11", 1)
        writeIpv4Rules(s2, "10.0.1.1", "08:00:00:00:01:11", 1)
        writeIpv4Rules(s1, "10.0.1.1", "08:00:00:00:01:11", 1)
        writeArpRules(s1, "10.0.1.1", "08:00:00:00:01:11", 2)
        writeArpRules(s1, "10.0.1.10", "08:00:00:00:01:00", 2)
        writeArpRules(s4, "10.0.2.2", "08:00:00:00:02:22", 2)
        writeArpRules(s4, "10.0.2.20", "08:00:00:00:02:00", 2)
        writeIpv6Rules(s1, "fe80::5678", 128, "08:00:00:00:02:22", 3)
        writeIpv6Rules(s3, "fe80::5678", 128, "08:00:00:00:02:22", 2)
        writeIpv6Rules(s4, "fe80::5678", 128, "08:00:00:00:02:22", 2)


        writeIcmpEncapsulationRule(s1,
                                    is_ipv4_encap=True,
                                    src_ipv4="10.0.1.1",
                                    dst_ipv4="10.0.2.2",
                                    outer_dst_mac="08:00:00:00:01:00",
                                    outer_src_ip="10.0.1.10",
                                    outer_dst_ip="10.0.2.20",
                                    port=2
                                    )

        # 示例 2：将 h1（10.0.1.1）发往 h2（10.0.2.2）的 ICMP 包封装为 IPv6
        # writeIcmpEncapsulationRule(
        #     s1,
        #     is_ipv4_encap=False,
        #     src_ipv4="10.0.1.1",
        #     dst_ipv4="10.0.2.2",
        #     outer_dst_mac="08:00:00:00:01:00",  # s1-p2 的 MAC
        #     outer_src_ip="fe80::1234",  # h1 的 IPv6 地址（需为交换机支持的 IPv6 地址）
        #     outer_dst_ip="fe80::5678",  # h2 的 IPv6 地址
        #     port=2  # 出端口为 s1-p2
        # )

        print("\nAll rules installed. Press Ctrl+C to stop.")
        while True:
            sleep(1)

    except KeyboardInterrupt:
        print("\nShutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    finally:
        ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format',
                        type=str, default='./build/basic.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file',
                        type=str, default='./build/basic.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        print(f"P4Info file not found: {args.p4info}")
        sys.exit(1)
    if not os.path.exists(args.bmv2_json):
        print(f"BMv2 JSON file not found: {args.bmv2_json}")
        sys.exit(1)

    main(args.p4info, args.bmv2_json)
